(*
** Implementacion de AES (Rijndael) en el lenguaje definido en Lang.v.
*)

Require Export Lang.
Require Import TypesImpl.
Open Scope Z_scope.
Open Scope lang_scope.

Set Implicit Arguments.

Definition word := Matrix Byte 4 1.

(*Number of columns*)
Parameter Nb : Integer.
(*Number of rounds*)
Parameter Nr : Integer.
Parameter Nk : Integer.
Parameter Nk__Nr_1 : Integer.

Definition round := id 2.
Definition i := id 3.
Definition j := id 4.
Definition k := id 5.
Definition aux := id 6.
Definition sig := id 7.
Definition x := id 8.
Definition y := id 9.
Definition b := id 10.
Definition row := id 11.
Definition col := id 12.
Definition baux := id 13.
Definition w := id 14.
Definition sbox := id 15.
Definition a0 := id 16.
Definition a1 := id 17.
Definition a2 := id 18.
Definition a3 := id 19.
Definition gmul2a0 := id 20.
Definition gmul3a0 := id 21.
Definition gmul2a1 := id 22.
Definition gmul3a1 := id 23.
Definition gmul2a2 := id 24.
Definition gmul3a2 := id 25.
Definition gmul2a3 := id 26.
Definition gmul3a3 := id 27.
Definition temp := id 28.
Definition auxint := id 29.
Definition rcon := id 30.


Definition sub_word(w sbox: Id) :=
  int i ::= 1;
  FOR i TO 4 DO
    byte b ::= w[i][1];
    byte row ::= b band 240;(*11110000*)
    byte row ::= row shr 4;
    byte col ::= b band 15;(*00001111*)
    byte aux ::= sbox[row][col];
    w[i][1] ::= aux;
  END
.

Definition rot_word(w: Id) :=
  int k ::= 1;
  byte aux ::= w[1][1];
  FOR k TO (Nb -1) DO
    byte sig ::= w[k + 1][1];
    w[k][1] ::= sig;
  END;
  w[4][1] ::= aux;
.

Definition key_expansion(key: Id):=
  int i ::= 1;
  WHILE (i le Nk) DO
    (* w[i] = word(key[4*i], key[4*i+1], key[4*i+2], key[4*i+3]) *)
    byte aux ::= key[1][i];
    w[1][i] ::= aux;

    byte aux ::= key[2][i];
    w[2][i] ::= aux;

    byte aux ::= key[3][i];
    w[3][i] ::= aux;

    byte aux ::= key[4][i];
    w[4][i] ::= aux;

    int i ::= i + 1;
  END;

  int i ::= Nk;
  int auxint ::= Nk * (Nr + 1);

  WHILE (i le auxint) DO
    byte aux ::= w[1][i - 1];
    temp[1][1] ::= aux;
    byte aux ::= w[2][i - 1];
    temp[1][2] ::= aux;
    byte aux ::= w[3][i - 1];
    temp[1][3] ::= aux;
    byte aux ::= w[4][i - 1];
    temp[1][4] ::= aux;

    IFB (i mod Nk == 1) THEN
      rot_word temp;
      sub_word temp sbox;

      byte temp ::= temp xor rcon[1][i / Nk];
    ELSE
      sub_word temp sbox;
    END;

    (* w[i] = w[i-Nk] xor temp *)
    byte aux ::= w[1][i - Nk] xor temp;
    w[i - 1][1] ::= aux;
    byte aux ::= w[2][i - Nk] xor temp;
    w[i - 1][2] ::= aux;
    byte aux ::= w[3][i - Nk] xor temp;
    w[i - 1][3] ::= aux;
    byte aux ::= w[4][i - Nk];
    w[i - 1][4] ::= aux;

    int i ::= i + 1;
  END
.


Definition sub_bytes(matr: Id)(sbox: Id):=
  int i ::= 1;
  int j ::= 1;

  FOR i TO Nb DO
    FOR j TO 4 DO
      byte b ::= matr[i][j];
      byte row ::= b band 240;(*11110000*)
      byte row ::= row shr 4;
      byte col ::= b band 15;(*00001111*)
      byte aux ::= sbox[row][col];
      matr[i][j] ::= aux;
    END
  END
.

Definition add_round_key(matr: Id)(key: Id)(offset: ArithExp) :=
  int i ::= offset + 1;
  int j ::= 1;

  FOR i TO Nb DO
    FOR j TO 4 DO
      byte x ::= matr[i][j];
      byte y ::= key[i][j];
      matr[i][j] ::= x xor y;
    END
  END
.

(* Galois mult 2 *)
Definition gmul2(b: Id) :=
  byte b ::= b shl 1;
  IFB (b gt 255) THEN
    byte b ::= b xor 27;
  END
.

(* Galois mult 3 *)
Definition gmul3(b: Id) :=
  byte baux ::= b;
  byte b ::= b shl 1;

  IFB (b gt 255) THEN
    byte b ::= b xor 27;
  END;
  byte b ::= b xor baux;
.

Definition mix_columns(matr: Id) :=
  int i ::= 1;

  FOR i TO Nb DO
    byte a0 ::= matr[i][1];
    byte a1 ::= matr[i][2];
    byte a2 ::= matr[i][3];
    byte a3 ::= matr[i][4];

    byte gmul2a0 ::= a0;
    byte gmul3a0 ::= a0;
    byte gmul2a1 ::= a1;
    byte gmul3a1 ::= a1;
    byte gmul2a2 ::= a2;
    byte gmul3a2 ::= a2;
    byte gmul2a3 ::= a3;
    byte gmul3a3 ::= a3;

    gmul2 gmul2a0;
    gmul3 gmul3a0;
    gmul2 gmul2a1;
    gmul3 gmul3a1;
    gmul2 gmul2a2;
    gmul3 gmul3a2;
    gmul2 gmul2a3;
    gmul3 gmul3a3;

    matr[i][1] ::= gmul2a0 xor gmul3a1 xor a2 xor a3;
    matr[i][2] ::= a0 xor gmul2a1 xor gmul3a2 xor a3;
    matr[i][3] ::= a0 xor a1 xor gmul2a2 xor gmul3a3;
    matr[i][4] ::= gmul3a0 xor a1 xor a2 xor gmul2a3;
  END
.

Definition shift_rows(matr: Id) :=
  int i ::= 1;
  int j ::= 1;

  FOR i TO Nb DO
    FOR j TO i DO
      byte aux ::= matr[1][j];

      FOR k TO 3 DO
        byte sig ::= matr[i + 1][k];
        matr[i][k] ::= sig;
      END;

      matr[i][j] ::= aux;
    END
  END
.  

Definition aes_cipher (state: Id) :=
  add_round_key state w 1;
  int round ::= 1;
  FOR round TO Nr DO
    sub_bytes state sbox;
    shift_rows state;
    mix_columns state;
    add_round_key state w (round * Nb);
  END;

  sub_bytes state sbox;
  shift_rows state;
  add_round_key state w (round * Nb);
.
